#include "StdAfx.h"
#include "Util.h"
#include <sddl.h>

namespace UTIL {

CStringW DecodeFromUTF8(const CStringA &strSource)
{
	CStringW strResult;

	int iCount = MultiByteToWideChar(CP_UTF8, 0, strSource, strSource.GetLength(), strResult.GetBuffer(strSource.GetLength()), strSource.GetLength()+1);
	strResult.ReleaseBuffer(iCount);

	return strResult;
}

CStringA EncodeToBASE64(const BYTE *pSource, int iLength)
{
	int nLength = Base64EncodeGetRequiredLength(iLength);

	CStringA strResult;
	if (!Base64Encode(pSource, iLength, strResult.GetBuffer(nLength), &nLength, ATL_BASE64_FLAG_NOCRLF))
	{
		ATLASSERT(FALSE);
	}

	strResult.ReleaseBuffer(nLength);

	return strResult;
}

CAtlArray<BYTE> &DecodeBASE64(const CStringA &strSource, CAtlArray<BYTE> &arResult)
{
	int nLength = Base64DecodeGetRequiredLength(strSource.GetLength());

	arResult.SetCount(nLength);
	if (!Base64Decode(strSource, strSource.GetLength(), arResult.GetData(), &nLength))
	{
		ATLASSERT(FALSE);
	}
	arResult.SetCount(nLength);

	return arResult;
}

bool GetVersion(LPCTSTR pszFile, VS_FIXEDFILEINFO &rInfo)
{
	DWORD dwHandle = 0;
	DWORD dwSize = ::GetFileVersionInfoSize(pszFile, &dwHandle);
	if (!dwSize) return false;

	CAtlArray<BYTE> buffer;
	if (!buffer.SetCount(dwSize))
		return false;

	ZeroMemory(buffer.GetData(), dwSize);
	if (!::GetFileVersionInfo(pszFile, dwHandle, dwSize, buffer.GetData()))
		return false;

	VS_FIXEDFILEINFO* info = NULL; UINT uint = 0;
	if (!::VerQueryValue(buffer.GetData(), _T("\\"), (LPVOID*)&info, &uint)) 
		return false;

	rInfo = *info;
	return true;
}

CString GetVersion()
{
	try
	{
		TCHAR this_file_name[MAX_PATH+1];
		ZeroMemory(this_file_name, sizeof(this_file_name));
		if (!GetModuleFileName(_AtlBaseModule.GetModuleInstance(), this_file_name, MAX_PATH))
			return CString();

		VS_FIXEDFILEINFO info;
		ZeroMemory(&info, sizeof(info));
		if (!GetVersion(this_file_name, info))
			return CString();

		CString strVersion;
		strVersion.Format( _T("%i.%i.%i.%i"), info.dwFileVersionMS>>16, info.dwFileVersionMS&0xFF, info.dwFileVersionLS>>16, info.dwFileVersionLS&0xFF );
		ATLTRACE( _T("Current File Version: %s\n"), strVersion );
		return strVersion;
	}
	catch(...)
	{
		ATLTRACE( _T("Exception on GetVersion()\n") );
	}

	return CString();
}

bool IsAttachedToIE7OrHigher()
{
	if (!IsIE())
		return false;

	try
	{
		TCHAR path[MAX_PATH+1];
		ZeroMemory(path, sizeof(path));
		if (!GetModuleFileName(NULL, path, MAX_PATH))
			return false;

		VS_FIXEDFILEINFO info;
		ZeroMemory(&info, sizeof(info));
		if (!GetVersion(path, info))
			return false;

		ATLTRACE( _T("Current IE Version: %i.%i.%i.%i\n"), info.dwFileVersionMS>>16, info.dwFileVersionMS&0xFF, info.dwFileVersionLS>>16, info.dwFileVersionLS&0xFF );
		return (info.dwFileVersionMS>>16) >= 7;
	}
	catch(...)
	{
		ATLTRACE( _T("Exception on IsAttachedToIE7OrHigher()\n") );
	}

	return false;
}

// Set Thread Integrity Level Low

typedef struct _TOKEN_MANDATORY_LABEL {
	SID_AND_ATTRIBUTES Label;
} TOKEN_MANDATORY_LABEL, *PTOKEN_MANDATORY_LABEL;

typedef WINADVAPI BOOL (WINAPI *CONVERTSTRINGSIDTOSID) (LPCTSTR, PSID);

#define SE_GROUP_INTEGRITY			(0x00000020L)
#define SE_GROUP_INTEGRITY_ENABLED	(0x00000040L)

HRESULT SetThreadIntegrityLevelLow(HANDLE *phThread)
{
	Handle hProcToken;
	if (!OpenProcessToken(GetCurrentProcess(),MAXIMUM_ALLOWED, &hProcToken))
		return HRESULT_FROM_WIN32(GetLastError());

	Handle hMICToken;
	if (!DuplicateTokenEx(hProcToken,
		MAXIMUM_ALLOWED,
		NULL,
		SecurityImpersonation,
		TokenImpersonation,
		&hMICToken))
		return HRESULT_FROM_WIN32(GetLastError());

	const TCHAR szIntegritySid[20] = _T("S-1-16-4096");

	PSID pMICSid = NULL;
	if (!ConvertStringSidToSid(szIntegritySid, &pMICSid))
		return HRESULT_FROM_WIN32(GetLastError());

	TOKEN_MANDATORY_LABEL TML = {0};
	//Set Process IL to Low
	TML.Label.Attributes = SE_GROUP_INTEGRITY | SE_GROUP_INTEGRITY_ENABLED;
	TML.Label.Sid = pMICSid;

	BOOL bSetToken = FALSE;
	if (SetTokenInformation(hMICToken,(TOKEN_INFORMATION_CLASS)25,&TML,sizeof(TML) + GetLengthSid(pMICSid)))
	{
		bSetToken = SetThreadToken(phThread, hMICToken);
	}

	LocalFree(pMICSid);

	return bSetToken ? S_OK : HRESULT_FROM_WIN32(GetLastError());
}


HRESULT SetThreadIntegrityLevelLow()
{
	HANDLE hThread = GetCurrentThread();

	if (!ImpersonateSelf(SecurityImpersonation))
		return HRESULT_FROM_WIN32(GetLastError());

	return SetThreadIntegrityLevelLow(&hThread);
}

BOOL CreateLowProcess(LPCTSTR pszApplicationName, LPTSTR pszCommandLine/* = NULL*/)
{
	ProcessEx process;

	return CreateLowProcess(process, pszApplicationName, pszCommandLine);
}

BOOL CreateLowProcess(ProcessEx &process, LPCTSTR pszApplicationName, LPTSTR pszCommandLine/* = NULL*/)
{
	Handle hProcToken;
	if (!OpenProcessToken(GetCurrentProcess(),MAXIMUM_ALLOWED, &hProcToken))
		return HRESULT_FROM_WIN32(GetLastError());

	Handle hNewToken;
	if (!DuplicateTokenEx(hProcToken, MAXIMUM_ALLOWED, NULL, SecurityImpersonation, TokenPrimary, &hNewToken))
		return HRESULT_FROM_WIN32(GetLastError());

	// Low integrity SID
	WCHAR wszIntegritySid[20] = L"S-1-16-4096";

	PSID pIntegritySid = NULL;
	if (!ConvertStringSidToSid(wszIntegritySid, &pIntegritySid))
		return HRESULT_FROM_WIN32(GetLastError());

	TOKEN_MANDATORY_LABEL TIL = {0};
	TIL.Label.Attributes = SE_GROUP_INTEGRITY;
	TIL.Label.Sid = pIntegritySid;

	BOOL bRet = FALSE;

	// Set the process integrity level
	if (SetTokenInformation(hNewToken, (TOKEN_INFORMATION_CLASS)25, &TIL, sizeof(TIL) + GetLengthSid(pIntegritySid)))
	{
		// Create the new process at Low integrity
		bRet = process.CreateAsUser(hNewToken, pszApplicationName, pszCommandLine);
	}

	LocalFree(pIntegritySid);

	return bRet ? S_OK : HRESULT_FROM_WIN32(GetLastError());
}

void ShowProcessIntegrityLevel(HANDLE hToken)
{
	DWORD dwLengthNeeded;
	// Get the Integrity level.
	if (!GetTokenInformation(hToken, (TOKEN_INFORMATION_CLASS)25, NULL, 0, &dwLengthNeeded))
	{
		DWORD dwError = GetLastError();
		if (dwError == ERROR_INSUFFICIENT_BUFFER)
		{
			PTOKEN_MANDATORY_LABEL pTIL = (PTOKEN_MANDATORY_LABEL)LocalAlloc(0, dwLengthNeeded);
			if (pTIL != NULL)
			{
				if (GetTokenInformation(hToken, (TOKEN_INFORMATION_CLASS)25, pTIL, dwLengthNeeded, &dwLengthNeeded))
				{
					DWORD dwIntegrityLevel = *GetSidSubAuthority(pTIL->Label.Sid, 
						(DWORD)(UCHAR)(*GetSidSubAuthorityCount(pTIL->Label.Sid)-1));

					if (dwIntegrityLevel < 0x00002000/*SECURITY_MANDATORY_MEDIUM_RID*/)
					{
						MessageBox(0, _T("Low"), _T("Integrity Level"), MB_OK);
					}
					else if (dwIntegrityLevel >= 0x00002000/*SECURITY_MANDATORY_MEDIUM_RID*/ && 
						dwIntegrityLevel < 0x00003000L/*SECURITY_MANDATORY_HIGH_RID*/)
					{
						MessageBox(0, _T("Medium"), _T("Integrity Level"), MB_OK);
					}
					else if (dwIntegrityLevel >= 0x00003000L/*SECURITY_MANDATORY_HIGH_RID*/)
					{
						MessageBox(0, _T("High"), _T("Integrity Level"), MB_OK);
					}
				}
				LocalFree(pTIL);
			}
		}
	}
}

void ShowProcessIntegrityLevel()
{
	Handle hProcToken;
	if (OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY | TOKEN_QUERY_SOURCE, &hProcToken))
		ShowProcessIntegrityLevel(hProcToken);
}

void ShowThreadIntegrityLevel()
{
	HANDLE hThread = GetCurrentThread();

	Handle hThreadToken;
	if (OpenThreadToken(hThread, TOKEN_QUERY | TOKEN_QUERY_SOURCE, FALSE, &hThreadToken))
		ShowProcessIntegrityLevel(hThreadToken);
}

}; //namespace UTIL
